#pragma once

#include <fstream>
#include <map>
#include <sstream>
#include <string>
#include <vector>

#include "NvInfer.h"
#include "NvInferRuntime.h"

using namespace nvinfer1;

namespace algo {

std::map<std::string, Weights> LoadWeights(const std::string file);

IScaleLayer *addBatchNorm2d(INetworkDefinition *network,
                            std::map<std::string, Weights> &weightMap,
                            ITensor &input, const std::string &lname,
                            float eps);

ILayer *convBlock(INetworkDefinition *network,
                  std::map<std::string, Weights> &weightMap, ITensor &input,
                  int outch, int ksize, int s, int g, const std::string &lname);

ILayer *focus(INetworkDefinition *network,
              std::map<std::string, Weights> &weightMap, ITensor &input,
              int inch, int outch, int ksize, const std::string &lname);

ILayer *bottleneck(INetworkDefinition *network,
                   std::map<std::string, Weights> &weightMap, ITensor &input,
                   int c1, int c2, bool shortcut, int g, float e,
                   const std::string &lname);

ILayer *C3(INetworkDefinition *network,
           std::map<std::string, Weights> &weightMap, ITensor &input, int c1,
           int c2, int n, bool shortcut, int g, float e,
           const std::string &lname);

ILayer *SPP(INetworkDefinition *network,
            std::map<std::string, Weights> &weightMap, ITensor &input, int c1,
            int c2, int k1, int k2, int k3, const std::string &lname);

std::vector<std::vector<float>> GetAnchors(
    std::map<std::string, Weights> &weightMap, const std::string &lname);

IPluginV2Layer *AddYoloLayer(INetworkDefinition *network,
                             std::map<std::string, Weights> &weightMap,
                             const std::string &lname,
                             std::vector<IConvolutionLayer *> dets);
}  // namespace algo
